#include <iostream>
#include <fstream>
#include <vector>
#include <unistd.h>
#include <opus/opus.h>
#include <ogg/ogg.h>

const int SAMPLE_RATE = 48000;
const int CHANNELS = 2;
const int FRAME_SIZE = 960;
const int APPLICATION = OPUS_APPLICATION_AUDIO;

// 编码函数
std::vector<unsigned char> encode_opus(const std::vector<opus_int16>& pcm_data)
{
    OpusEncoder* encoder;
    int error;

    // 初始化编码器
    encoder = opus_encoder_create(SAMPLE_RATE, CHANNELS, APPLICATION, &error);
    if (error != OPUS_OK)
    {
        std::cerr << "Failed to create Opus encoder: " << opus_strerror(error) << std::endl;
        return {};
    }

    // 设置比特率
    error = opus_encoder_ctl(encoder, OPUS_SET_BITRATE(128000));
    if (error != OPUS_OK)
    {
        std::cerr << "Failed to set bitrate: " << opus_strerror(error) << std::endl;
        opus_encoder_destroy(encoder);
        return {};
    }

    std::vector<unsigned char> encoded_data(pcm_data.size() * 2);
    int encoded_bytes = opus_encode(encoder, pcm_data.data(), FRAME_SIZE, encoded_data.data(), encoded_data.size());
    if (encoded_bytes < 0)
    {
        std::cerr << "Encoding failed: " << opus_strerror(encoded_bytes) << std::endl;
    }
    else
    {
        encoded_data.resize(encoded_bytes);
    }

    // 销毁编码器
    opus_encoder_destroy(encoder);

    return encoded_data;
}

// 读取 WAV 文件的 PCM 数据
std::vector<opus_int16> read_wav_file(const std::string& filename)
{
    std::ifstream file(filename, std::ios::binary);
    if (!file)
    {
        std::cerr << "Failed to open WAV file: " << filename << std::endl;
        return {};
    }

    // 跳过 WAV 文件头（通常 44 字节）
    file.seekg(44, std::ios::beg);

    std::vector<opus_int16> pcm_data;
    opus_int16 sample;
    while (file.read(reinterpret_cast<char*>(&sample), sizeof(opus_int16)))
    {
        pcm_data.push_back(sample);
    }

    return pcm_data;
}

// 解码函数
std::vector<opus_int16> decode_opus(const std::vector<unsigned char>& opus_data)
{
    OpusDecoder* decoder;
    int error;

    // 初始化解码器
    decoder = opus_decoder_create(SAMPLE_RATE, CHANNELS, &error);
    if (error != OPUS_OK)
    {
        std::cerr << "Failed to create Opus decoder: " << opus_strerror(error) << std::endl;
        return {};
    }

    std::vector<opus_int16> decoded_data(FRAME_SIZE * CHANNELS);
    int decoded_samples = opus_decode(decoder, opus_data.data(), static_cast<int>(opus_data.size()), decoded_data.data(), FRAME_SIZE, 0);
    if (decoded_samples < 0)
    {
        std::cerr << "Decoding failed: " << opus_strerror(decoded_samples) << std::endl;
    }
    else
    {
        decoded_data.resize(decoded_samples * CHANNELS);
    }

    // 销毁解码器
    opus_decoder_destroy(decoder);

    return decoded_data;
}

// 写入 WAV 文件头
void write_wav_header(std::ofstream& file, int sampleRate, int channels, int bitsPerSample, int dataSize)
{
    int chunkSize = 36 + dataSize;
    int subchunk1Size = 16;
    int audioFormat = 1;
    int byteRate = sampleRate * channels * (bitsPerSample / 8);
    int blockAlign = channels * (bitsPerSample / 8);

    file.write("RIFF", 4);
    file.write(reinterpret_cast<const char*>(&chunkSize), 4);
    file.write("WAVE", 4);
    file.write("fmt ", 4);
    file.write(reinterpret_cast<const char*>(&subchunk1Size), 4);
    file.write(reinterpret_cast<const char*>(&audioFormat), 2);
    file.write(reinterpret_cast<const char*>(&channels), 2);
    file.write(reinterpret_cast<const char*>(&sampleRate), 4);
    file.write(reinterpret_cast<const char*>(&byteRate), 4);
    file.write(reinterpret_cast<const char*>(&blockAlign), 2);
    file.write(reinterpret_cast<const char*>(&bitsPerSample), 2);
    file.write("data", 4);
    file.write(reinterpret_cast<const char*>(&dataSize), 4);
}

// 编码 WAV 文件到 OPUS 文件并封装为 Ogg
void encode_wav_to_opus(const std::string& input_wav_filename, const std::string& output_opus_filename)
{
    std::vector<opus_int16> pcm_data = read_wav_file(input_wav_filename);
    if (pcm_data.empty())
    {
        return;
    }

    ogg_stream_state os;
    ogg_page og;
    ogg_packet op;
    int serialno = 0;
    int packetno = 0;

    ogg_stream_init(&os, serialno);

    std::ofstream output_opus_file(output_opus_filename, std::ios::binary);
    if (!output_opus_file)
    {
        std::cerr << "Failed to open output OPUS file: " << output_opus_filename << std::endl;
        ogg_stream_clear(&os);
        return;
    }

    // 按帧处理 PCM 数据并编码
    for (size_t i = 0; i < pcm_data.size(); i += FRAME_SIZE * CHANNELS)
    {
        std::vector<opus_int16> frame_data;
        if (i + FRAME_SIZE * CHANNELS <= pcm_data.size())
        {
            frame_data.assign(pcm_data.begin() + i, pcm_data.begin() + i + FRAME_SIZE * CHANNELS);
        }
        else
        {
            // 处理最后一帧不足的情况
            frame_data.assign(pcm_data.begin() + i, pcm_data.end());
            // 填充到一帧的大小
            while (frame_data.size() < FRAME_SIZE * CHANNELS)
            {
                frame_data.push_back(0);
            }
        }

        std::vector<unsigned char> encoded_frame = encode_opus(frame_data);

        if (!encoded_frame.empty())
        {
            op.bytes = encoded_frame.size();
            op.packet = encoded_frame.data();
            op.b_o_s = (i == 0);
            op.e_o_s = (i + FRAME_SIZE * CHANNELS >= pcm_data.size());
            op.granulepos = i / (FRAME_SIZE * CHANNELS);
            op.packetno = packetno++;

            ogg_stream_packetin(&os, &op);

            while (ogg_stream_pageout(&os, &og) == 1)
            {
                output_opus_file.write(reinterpret_cast<const char*>(og.header), og.header_len);
                output_opus_file.write(reinterpret_cast<const char*>(og.body), og.body_len);
            }
        }
    }

    ogg_stream_clear(&os);
    output_opus_file.close();
    std::cout << "Encoded data saved to " << output_opus_filename << std::endl;
}

// 解码 OPUS 文件到 WAV 文件
void decode_opus_to_wav(const std::string& input_opus_filename, const std::string& output_wav_filename)
{
    std::ifstream input_opus_file(input_opus_filename, std::ios::binary);
    if (!input_opus_file)
    {
        std::cerr << "Failed to open input OPUS file: " << input_opus_filename << std::endl;
        return;
    }

    ogg_sync_state oy;
    ogg_page og;
    ogg_packet op;

    ogg_sync_init(&oy);

    std::vector<unsigned char> opus_data;
    char* buffer = new char[4096];
    int buffer_size = 4096;

    while (input_opus_file.read(buffer, buffer_size))
    {
        int bytes_read = input_opus_file.gcount();
        ogg_sync_wrote(&oy, bytes_read);
        while (ogg_sync_pageout(&oy, &og) == 1)
        {
            ogg_stream_state os;
            ogg_stream_init(&os, ogg_page_serialno(&og));

            int result = ogg_page_bos(&og);
            if (result)
            {
                ogg_stream_reset(&os);
            }

            while (ogg_stream_packetout(&os, &op) == 1)
            {
                std::vector<unsigned char> packet_data(op.packet, op.packet + op.bytes);
                opus_data.insert(opus_data.end(), packet_data.begin(), packet_data.end());
            }

            ogg_stream_clear(&os);
        }
    }

    delete[] buffer;
    ogg_sync_clear(&oy);
    input_opus_file.close();

    // 解码 OPUS 数据
    std::vector<opus_int16> decoded_pcm_data = decode_opus(opus_data);

    // 写入解码后的 PCM 数据到 WAV 文件
    std::ofstream output_wav_file(output_wav_filename, std::ios::binary);
    if (!output_wav_file)
    {
        std::cerr << "Failed to open output decoded WAV file: " << output_wav_filename << std::endl;
        return;
    }

    int dataSize = static_cast<int>(decoded_pcm_data.size()) * sizeof(opus_int16);
    write_wav_header(output_wav_file, SAMPLE_RATE, CHANNELS, 16, dataSize);
    output_wav_file.write(reinterpret_cast<const char*>(decoded_pcm_data.data()), dataSize);
    output_wav_file.close();

    std::cout << "Decoded data saved to " << output_wav_filename << std::endl;
}

int main(int argc, char* argv[])
{
    if (argc != 4)
    {
        std::cerr << "Usage: " << argv[0] << " <encode|decode> <input_file> <output_file>" << std::endl;
        return 1;
    }

    std::string action = argv[1];
    std::string input_file = argv[2];
    std::string output_file = argv[3];

    if (action == "encode")
    {
        encode_wav_to_opus(input_file, output_file);
    }
    else if (action == "decode")
    {
        decode_opus_to_wav(input_file, output_file);
    }
    else
    {
        std::cerr << "Invalid action. Use 'encode' or 'decode'." << std::endl;
        return 1;
    }

    return 0;
}